import pulp
import numpy as np
from pprint import pprint

if __name__ == '__main__':
    costs = np.array([[100,700,200,0,0],
                       [0,200,800,0,0],
                       [100,100,100,100,600],
                       [871,0,99,0,30],
                       [0,33,33,34,900],
                       [250,250,0,250,250],
                       [111,1,0,333,555]])

    max_plant = [1,1,1,1,1,1,1]
    max_cultivation = [1,1,1,1,1,2,2]

def transportation_problem(costs, x_max, y_max):

    row = len(costs)
    col = len(costs[0])

    prob = pulp.LpProblem('Transportation Problem', sense=pulp.LpMaximize)

    var = [[pulp.LpVariable(f'x{i}{j}', lowBound=0, cat=pulp.LpInteger) for j in range(col)] for i in range(row)]

    flatten = lambda x: [y for l in x for y in flatten(l)] if type(x) is list else [x]

    prob += pulp.lpDot(flatten(var), costs.flatten())

    for i in range(row):
        prob += (pulp.lpSum(var[i]) <= x_max[i])

    for j in range(col):
        prob += (pulp.lpSum([var[i][j] for i in range(row)]) <= y_max[j])

    prob.solve()

    return {'objective':pulp.value(prob.objective), 'var': [[pulp.value(var[i][j]) for j in range(col)] for i in range(row)]}
    
res = transportation_problem(costs, max_plant, max_cultivation)
print(f'最大值为{res["objective"]}')
print('各变量的取值为：')
pprint(res['var'])    